import os
import rospy
from sensor_msgs.msg import Range
import pandas as pd
from PIL import Image
import cv2
import numpy as np

# 默认 f_dx 值
DEFAULT_F_DX = 1.5

def get_distance_from_rostopic(topic_name):
    """
    从指定的 rostopic 读取距离数据。
    """
    # 初始化 ROS 节点
    rospy.init_node('distance_listener', anonymous=True)
    
    # 订阅指定的 topic
    distance = None

    def callback(msg):
        nonlocal distance
        distance = msg.range

    rospy.Subscriber(topic_name, Range, callback)
    
    # 等待距离数据的回调
    rospy.sleep(1)
    
    return distance

def rust_estimation(img_files_path, label_files_path, f_dx):
    img_files = os.listdir(img_files_path)
    result = []
    for img_name in img_files:
        img_path = os.path.join(img_files_path, img_name)
        label_path_root = os.path.join(label_files_path, img_name.split('.')[0])
        label_path = label_path_root + '.txt'
        img = Image.open(img_path)
        width = float(img.size[0])
        height = float(img.size[1])
        fo1 = open(label_path, 'r')
        lines = [l.split() for l in fo1.readlines() if l.strip()]
        S_uv = 0
        for line in lines:
            if line[0] == '2':
                # 像素数累加
                left = int(round(((float(line[1]) - 0.5*float(line[3]))*width), 0))
                top = int(round(((float(line[2]) - 0.5*float(line[4]))*height), 0))
                right = int(round(((float(line[1]) + 0.5*float(line[3]))*width), 0))
                low = int(round(((float(line[2]) + 0.5*float(line[4]))*height), 0))
                add_S_uv = num_uv(img_path, left, top, right, low)
                S_uv = S_uv + add_S_uv
            else:
                continue
        S_e = S_uv * f_dx * f_dx
        S_e = str(S_e) + ' mm^2'
        result.append([img_name, S_e])

    return result

def num_uv(img_path, left, top, right, low):
    mat_img = cv2.imread(img_path)
    cropImg = mat_img[top:low, left:right]

    # 提取中心点的像素rgb值
    x_center = int(round(0.5 * (right - left), 0))
    y_center = int(round(0.5 * (low - top), 0))
    rgb = cropImg[y_center][x_center]
    rgb_yuzhi = [x + 50 for x in rgb]
    num_uv_ = 0
    for i in range(cropImg.shape[0]):
        for j in range(cropImg.shape[1]):
            if (cropImg[i][j][0] < rgb_yuzhi[0]) & (cropImg[i][j][1] < rgb_yuzhi[1]) & (cropImg[i][j][2] < rgb_yuzhi[2]):
                num_uv_ = num_uv_ + 1
            else:
                cropImg[i][j] = [1, 1, 1]

    return num_uv_

if __name__ == "__main__":
    topic_name = "/distance_topic"  # 替换为你的实际 topic 名称
    distance = get_distance_from_rostopic(topic_name)
    if distance is not None:
        # 计算 f_dx 值
        f_dx = DEFAULT_F_DX / distance
        print(f"Calculated f_dx: {f_dx}")

        img_files_path_5 = "./runs/detect/exp_0.5m_new/images"
        label_files_path_5 = "./runs/detect/exp_0.5m_new/labels"
        img_files_path_1 = "./runs/detect/exp_1m_new/images"
        label_files_path_1 = "./runs/detect/exp_1m_new/labels"

        result_5 = rust_estimation(img_files_path_5, label_files_path_5, f_dx)
        result_1 = rust_estimation(img_files_path_1, label_files_path_1, f_dx)

        # 保存结果到文件
        with open("result_0.5m", 'a') as f:
            for line in result_5:
                f.writelines(line[0] + ': ' + line[1] + '\n')

        with open("result_1m", 'a') as f:
            for line in result_1:
                f.writelines(line[0] + ': ' + line[1] + '\n')

    else:
        print("未能从 rostopic 获取距离数据。")

